Summary

This PR introduces an end-to-end implementation of Step3.5 Multi-Token Prediction (MTP) in llama.cpp, covering model conversion, loading, runtime graph execution, speculative decoding, server integration, and quantization compatibility.

The current implementation is intentionally specialized for Step3.5 and focuses on enabling a complete and usable MTP pipeline within the existing architecture. It allows early experimentation with Step3.5-style MTP and provides a concrete reference for evaluating performance and behavior in practice.

This PR does not attempt to define a generalized MTP abstraction. Instead, it provides a concrete baseline that can be used for experimentation, benchmarking, and informing future design directions toward a more general solution.

Implemented

• Conversion and tensor mapping - export the Step3.5 nextn / shared-head tensors needed by the current runtime path and register them as proper per-layer tensors, currently keeping the first MTP layer only.
• Model loading - load the Step3.5 MTP weights when -mtp is enabled instead of always pulling the extra tensors into the runtime.
• Runtime graph - add a dedicated Step3.5 MTP graph path in the iswa builder, while refactoring layer construction so the normal decoder path and the MTP path share the same core logic.
• Speculation - add an MTP-specific speculative decoding state machine with retained prefix handling, staged first-pass sources, and verifier hidden-state handoff.
• SWA guard - add SWA / iSWA guard and rollback-related handling to preserve correct MTP behavior across SWA window boundaries, especially when advancing the draft would evict cache state that cannot be recovered after accept.
• Server path - wire prompt hidden-state collection, MTP draft setup, accept / rollback flow, and prefix reuse into the server implementation.
• CLI compatibility - expose the feature through CLI options and keep the speculative example usable without a separate draft model.
• Quantization compatibility - make the quantization path recognize the new MTP tensors and avoid incompatible imatrix requirements for them.

Not Yet Completed

• Multi-layer MTP - the current Step3.5 runtime only uses the first MTP layer.
• Generalization - the implementation is still Step3.5-oriented and is not yet shaped into a more general MTP framework.
• Cache reuse - only continuous prefix reuse is supported for MTP right now; the prompt-cache reuse path is currently disabled, and the more general cache reuse path is not handled yet.
• Shared draft context - the server still uses one draft context per slot and has not been reworked into a shared / batched draft-context design.

Testing

These numbers are taken from the server log line:

eval time = ... ms / ... tokens (... ms per token, ... tokens per second)

For these runs, -mtp enables the Step3.5 MTP path introduced in this PR: it loads the MTP weights and makes the server run the MTP speculative decoding flow being measured below.

Decode Speed on 8xH200

nohup env ./llama-server \
  -m ../../../step3p5_flash_fp16.gguf \
  --host 0.0.0.0 --port 8080 \
  --device CUDA0,CUDA1,CUDA2,CUDA3,CUDA4,CUDA5,CUDA6,CUDA7 \
  --split-mode layer --tensor-split 1,1,1,1,1,1,1,1 \
  -c 262144 -ctk f16 -ctv f16 \
  -np 1 -cb -b 4096 -ub 2048 \
  --reasoning-format none \
  -mtp --draft 3 -cram 0 \
  > ./llama-server_mtp.log 2>&1 &

Model weights: fp16, KV cache: fp16

Model weights: fp16, KV cache: q8_0

Model weights: fp16, KV cache: q4_0

Model weights: Q4_K_S, KV cache: fp16

Model weights: IQ4_XS, KV cache: fp16

Decode Speed on Mac Studio

Performance on Mac Studio is currently weaker than I expected. Profiling so far suggests that the main-model verify path on the Metal backend is relatively expensive. If anyone has suggestions or useful leads here, feedback would be very welcome.

nohup ./llama-server \
  -m ../../../step3p5_flash_IQ4_XS.gguf \
  --host 0.0.0.0 --port 8080 \
  -c 131072 \
  -ctk f16 -ctv f16 \
  -np 1 -cb \
  -b 4096 -ub 2048 \
  --jinja \
  --reasoning-format none \
  -mtp --draft 3 -cram 0 \
    > ./llama-server_mtp.log 2>&1 &

Model weights: IQ4_XS, KV cache: fp16

Model weights: Q4_K_S, KV cache: fp16, MTP KV cache: q4_0

Decode Speed on DGX Spark

nohup env GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 ./llama-server \
  --mlock \
  -m ../../../step3p5_flash_mtp_Q4_K_S.gguf \
  --host 0.0.0.0 --port 8080 \
  -c 131072 \
  -ctk q8_0 -ctv q8_0 \
  -ctkd q4_0 -ctvd q4_0 \
  -np 1 -cb \
  -b 2048 -ub 1024 \
  --jinja \
  --reasoning-format none \
  -mtp --draft 3 -cram 0 \
  > ./llama-server_mtp.log 2>&1 &

Model weights: IQ4_XS, KV cache: fp16

Model weights: Q4_K_S, KV cache: q8_0, MTP KV cache: q4_0

Related Work

• #18886
• #18039
• #15225